{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "absolute-birth",
   "metadata": {},
   "source": [
    "# MetaSpore Getting Started"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "personal-technique",
   "metadata": {
    "tags": []
   },
   "source": [
    "MetaSpore is a machine learning platform, which provides a one-stop solution for data preprocessing, model training and online prediction."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ignored-radiation",
   "metadata": {},
   "source": [
    "In this article, we introduce the basic API of MetaSpore briefly."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "welcome-bullet",
   "metadata": {},
   "source": [
    "## Prepare Data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "answering-coordinator",
   "metadata": {},
   "source": [
    "We use the publicly available dataset [Terabyte Click Logs](https://labs.criteo.com/2013/12/download-terabyte-click-logs-2/) published by CriteoLabs as our demo dataset."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "pretty-setup",
   "metadata": {},
   "source": [
    "We sample the dataset with sampling rate 0.001 so that the running of the demo can finish quickly. More information about the demo dataset can be found in [MetaSpore Demo Dataset](https://ks3-cn-beijing.ksyuncs.com/dmetasoul-bucket/demo/criteo/index.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "arabic-sheep",
   "metadata": {},
   "source": [
    "Execute the following cell to download the demo dataset into the working directory. Those data files take up about 2.1 GiB disk space and the downloading process may take sveral minutes. If the downloading fails, please refer to [MetaSpore Demo Dataset](https://ks3-cn-beijing.ksyuncs.com/dmetasoul-bucket/demo/criteo/index.html) and download the dataset manually."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "turkish-lebanon",
   "metadata": {},
   "outputs": [],
   "source": [
    "import metaspore\n",
    "metaspore.demo.download_dataset()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "awful-preference",
   "metadata": {},
   "source": [
    "You can check the downloaded dataset by executing the following cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "inclusive-cover",
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -l ${PWD}/data/"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "deluxe-plymouth",
   "metadata": {},
   "source": [
    "(Optional) To upload the dataset to your own s3 bucket:\n",
    "\n",
    "1. Fill ``{YOUR_S3_BUCKET}`` and ``{YOUR_S3_PATH}`` with your preferred values in the following cell.\n",
    "2. Uncomment the cell by removing the leading ``#`` character.\n",
    "3. Execute the cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "tired-plumbing",
   "metadata": {},
   "outputs": [],
   "source": [
    "#!aws s3 cp --recursive ${PWD}/data/ s3://{YOUR_S3_BUCKET}/{YOUR_S3_PATH}/demo/data/"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "naughty-programmer",
   "metadata": {},
   "source": [
    "Alternatively, you can open a terminal by selecting the ``File`` -> ``New`` -> ``Terminal`` menu item and executing Bash commands in it."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "false-guide",
   "metadata": {},
   "source": [
    "You can check the uploaded dataset in your s3 bucket by uncommenting and executing the following cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "little-password",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#!aws s3 ls s3://{YOUR_S3_BUCKET}/{YOUR_S3_PATH}/demo/data/"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f33805b-a786-4092-b280-e60a84308672",
   "metadata": {},
   "source": [
    "The ``schema`` directory contains configuration files and must also be uploaded to s3 so that the model can be trained in cluster environment. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82290030-1bf5-4e9a-8407-9972466c5816",
   "metadata": {},
   "outputs": [],
   "source": [
    "#!aws s3 cp --recursive ${PWD}/schema/ s3://{YOUR_S3_BUCKET}/{YOUR_S3_PATH}/demo/schema/"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ccbf941-347c-44a4-98b1-c0b2c831dd56",
   "metadata": {},
   "source": [
    "In the rest of the article, we assume the demo dataset and schemas has been uploaded to `ROOT_DIR`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbae0f69-e2e5-4bae-bcb2-dd6c74016a2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ROOT_DIR = 's3://{YOUR_S3_BUCKET}/{YOUR_S3_PATH}/demo'\n",
    "ROOT_DIR = '.'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "average-trustee",
   "metadata": {},
   "source": [
    "## Define the Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "equal-sacramento",
   "metadata": {},
   "source": [
    "We can define our neural network model by subclassing ``torch.nn.Module`` as usual PyTorch models. The following ``DemoModule`` class provides an example."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "later-insight",
   "metadata": {},
   "source": [
    "Compared to usual PyTorch models, the notable difference is the ``_sparse`` layer created by instantiating ``ms.EmbeddingSumConcat`` which takes an embedding size and paths of two text files. ``ms.EmbeddingSumConcat`` makes it possible to define large-scale sparse models in PyTorch, which is a distinguishing feature of MetaSpore."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7f30e2a-2fc3-4781-9927-300d74d260e2",
   "metadata": {},
   "source": [
    "The ``_schema_dir`` field is an s3 directory which makes it possible to use the ``DemoModule`` class in cluster environment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "korean-pierce",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import metaspore as ms\n",
    "\n",
    "class DemoModule(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self._embedding_size = 16\n",
    "        self._schema_dir = ROOT_DIR + '/schema/'\n",
    "        self._column_name_path = self._schema_dir + 'column_name_demo.txt'\n",
    "        self._combine_schema_path = self._schema_dir + 'combine_schema_demo.txt'\n",
    "        self._sparse = ms.EmbeddingSumConcat(self._embedding_size, self._column_name_path, self._combine_schema_path)\n",
    "        self._sparse.updater = ms.FTRLTensorUpdater()\n",
    "        self._sparse.initializer = ms.NormalTensorInitializer(var=0.01)\n",
    "        self._dense = torch.nn.Sequential(\n",
    "            ms.nn.Normalization(self._sparse.feature_count * self._embedding_size),\n",
    "            torch.nn.Linear(self._sparse.feature_count * self._embedding_size, 1024),\n",
    "            torch.nn.ReLU(),\n",
    "            torch.nn.Linear(1024, 512),\n",
    "            torch.nn.ReLU(),\n",
    "            torch.nn.Linear(512, 1),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self._sparse(x)\n",
    "        x = self._dense(x)\n",
    "        return torch.sigmoid(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "tested-vampire",
   "metadata": {},
   "source": [
    "Instantiating the ``DemoModule`` class to define our PyTorch model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "resistant-justice",
   "metadata": {},
   "outputs": [],
   "source": [
    "module = DemoModule()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dominant-intersection",
   "metadata": {},
   "source": [
    "## Train the Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "vocal-hollow",
   "metadata": {},
   "source": [
    "To train our model, first we need to create a ``ms.PyTorchEstimator`` passing in several arguments including our PyTorch model ``module`` and the number of workers and servers.\n",
    "\n",
    "``model_out_path`` specifies where to store the trained model.\n",
    "\n",
    "``input_label_column_index`` specifies the column index of the label column in the dataset, which is ``0`` for the demo dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "matched-realtor",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_out_path = ROOT_DIR + '/output/dev/model_out/'\n",
    "estimator = ms.PyTorchEstimator(module=module,\n",
    "                                worker_count=1,\n",
    "                                server_count=1,\n",
    "                                model_out_path=model_out_path,\n",
    "                                experiment_name='0.1',\n",
    "                                input_label_column_index=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "objective-hybrid",
   "metadata": {},
   "source": [
    "Next, we create a Spark session by calling ``ms.spark.get_session()`` and load the training dataset by call ``ms.input.read_s3_csv()``.\n",
    "\n",
    "``delimiter`` specifies the column delimiter of the dataset, which is the TAB character ``'\\t'`` for the demo dataset.\n",
    "\n",
    "We also need to pass column names because the csv files do not contain headers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a8d6c9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "column_names = []\n",
    "with open(f'{ROOT_DIR}/schema/column_name_demo.txt', 'r') as f:\n",
    "    for line in f:\n",
    "        column_names.append(line.split(' ')[1].strip())\n",
    "print(column_names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "friendly-launch",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset_path = ROOT_DIR + '/data/train/day_0_0.001_train.csv'\n",
    "spark_session = ms.spark.get_session(local=True,\n",
    "                                     batch_size=100,\n",
    "                                     worker_count=estimator.worker_count,\n",
    "                                     server_count=estimator.server_count,\n",
    "                                     log_level='INFO',\n",
    "                                     spark_confs={'spark.eventLog.enabled':'true'})\n",
    "train_dataset = ms.input.read_s3_csv(spark_session, train_dataset_path, delimiter='\\t', column_names=column_names)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "exact-hollow",
   "metadata": {},
   "source": [
    "Finally, we call the ``fit()`` method of ``ms.PyTorchEstimator`` to train our model. This will take several minutes and you can see the progress by looking at the output of the cell. The trained model is stored in ``model_out_path`` and the ``model`` variable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "higher-split",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = estimator.fit(train_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "rotary-quest",
   "metadata": {},
   "source": [
    "## Evaluate the Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "perfect-debut",
   "metadata": {},
   "source": [
    "To evaluate our model, we use the ``ms.input.read_s3_csv()`` function again to load the test dataset, passing in the column delimiter ``'\\t'``."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "infectious-anthropology",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_dataset_path = ROOT_DIR + '/data/test/day_0_0.001_test.csv'\n",
    "test_dataset = ms.input.read_s3_csv(spark_session, test_dataset_path, delimiter='\\t', column_names=column_names)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "extensive-timeline",
   "metadata": {},
   "source": [
    "Next, we call the ``model.transform()`` method to transform the test dataset, which will add a column named ``rawPrediction`` to the test dataset representing the predicted labels. For ease of integration with Spark MLlib, ``model.transform()`` will also add a column named ``label`` to the test dataset representing the actual labels.\n",
    "\n",
    "Like the training process, this will take several minutes and you can see the progress by looking at the output of the cell. The transformed test dataset is stored in the ``result`` variable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "exterior-pilot",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = model.transform(test_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "familiar-yesterday",
   "metadata": {},
   "source": [
    "``result`` is a normal PySpark DataFrame and can be inspected by its methods."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cognitive-illinois",
   "metadata": {},
   "outputs": [],
   "source": [
    "result.show(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "immediate-rebate",
   "metadata": {},
   "source": [
    "Finally, we use ``pyspark.ml.evaluation.BinaryClassificationEvaluator`` to compute test AUC."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "geographic-football",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import pyspark\n",
    "evaluator = pyspark.ml.evaluation.BinaryClassificationEvaluator()\n",
    "test_auc = evaluator.evaluate(result)\n",
    "print('test_auc: %g' % test_auc)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "256721e7-144d-479a-8fbd-0be3479137a2",
   "metadata": {},
   "source": [
    "When all computations are done, we should call the ``stop()`` method of ``spark_session`` to make sure all the resources are released."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb448788-b6ad-466c-8e9d-ff38d5235dab",
   "metadata": {},
   "outputs": [],
   "source": [
    "spark_session.stop()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "different-company",
   "metadata": {},
   "source": [
    "## Summary"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "juvenile-outside",
   "metadata": {},
   "source": [
    "We illustrated how to train and evaluate neural network model in MetaSpore. Users familiar with PyTorch and Spark MLlib should get started easily, which is the design goal of MetaSpore."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
